///////////////////////////////////////////////////////////////////////////////
///
/// Copyright (c) 2011 - PG-soft
///
/// Original filename: vmxdbg.c
/// Project          : vmxdbg
/// Date of creation : 2011-09-06
/// Author(s)        : ProgrammeBoy
///
/// Purpose          : <description>
///
/// Revisions:
///  0000 [2011-09-06] Initial revision.
///
///////////////////////////////////////////////////////////////////////////////

// $Id$

#include <ntddk.h>
#include <string.h>

#include "vt.h"
#include "cpu.h"
#include "vmx.h"
#include "gui.h"
#include "vmxdbg.h"

PDRIVER_OBJECT pdoGlobalDrvObj = 0;

NTSTATUS VMXDBG_DispatchCreateClose(
    IN PDEVICE_OBJECT		DeviceObject,
    IN PIRP					Irp
    )
{
    NTSTATUS status = STATUS_SUCCESS;
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return status;
}

NTSTATUS VMXDBG_DispatchDeviceControl(
    IN PDEVICE_OBJECT		DeviceObject,
    IN PIRP					Irp
    )
{
    NTSTATUS status = STATUS_SUCCESS;
    PIO_STACK_LOCATION irpSp = IoGetCurrentIrpStackLocation(Irp);

    switch(irpSp->Parameters.DeviceIoControl.IoControlCode)
    {
    default:
        Irp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;
        Irp->IoStatus.Information = 0;
        break;
    }

    status = Irp->IoStatus.Status;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return status;
}

VOID VMXDBG_DriverUnload(
    IN PDRIVER_OBJECT		DriverObject
    )
{
    PDEVICE_OBJECT pdoNextDeviceObj = pdoGlobalDrvObj->DeviceObject;
    ULONG i;
	KIRQL OldIrql;

	IoDeleteSymbolicLink(&usSymlinkName);

    //
	// Delete all the device objects
    //
	
	while(pdoNextDeviceObj)
    {
        PDEVICE_OBJECT pdoThisDeviceObj = pdoNextDeviceObj;
        pdoNextDeviceObj = pdoThisDeviceObj->NextDevice;
        IoDeleteDevice(pdoThisDeviceObj);
    }
	
	//
	// STOP VT
	//
	
	KdPrint(("Stop VT...\n"));
	for (i = 0; i < KeNumberProcessors; i++){
		KeSetSystemAffinityThread((KAFFINITY) (1 << i));
		OldIrql = KeRaiseIrqlToDpcLevel();
		_StopVT();
		KeLowerIrql(OldIrql);
	}

}

VOID 
KeSleep(
        ULONG milliSeconds
    )
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    PKTIMER timer = (PKTIMER)ExAllocatePoolWithTag( NonPagedPool, 
                                sizeof (KTIMER), 'HELO' );
    LARGE_INTEGER duetime;

    if (!timer)
        return;

    duetime.QuadPart = (ULONG64) milliSeconds * -10000;
    KeInitializeTimerEx(timer, NotificationTimer);
    KeSetTimerEx(timer, duetime, 0, NULL);
    KeWaitForSingleObject (timer, Executive, KernelMode, FALSE, NULL);
    ExFreePoolWithTag(timer, 'HELO');
}

NTSTATUS 
StartVT(
        PVOID GuestEsp
    )
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    NTSTATUS Status;
    PCPU_VM_CONTEXT pCPU;

    pCPU = (PCPU_VM_CONTEXT)ExAllocatePoolWithTag(NonPagedPool, 
            sizeof(CPU_VM_CONTEXT), 'neek');
    RtlZeroMemory( pCPU, sizeof(CPU_VM_CONTEXT) );
    if (!pCPU){
        KdPrint(( "can't allocate cpu structure\n" ));
        return STATUS_INSUFFICIENT_RESOURCES;
    }
    
    //
    //  Get Current Processor number
    //
    
    pCPU->CpuIndex = KeGetCurrentProcessorNumber();
    
	if ( !CheckIfVMXIsEnabled() ){
		KdPrint(("Please Enable VT in BIOS\n"));
		return STATUS_SUCCESS;
	}

    //
    //  Set VMXON
    //
    
    if ( !NT_SUCCESS(InitializeVMX(pCPU)) ){
        KdPrint(("InitializeVMX Error!\n"));
        return STATUS_UNSUCCESSFUL;
    }

    KdPrint(("InitializeVMX ok!\n"));
    
    //
    //  Set VMCS
    //
    
    if ( !NT_SUCCESS(InitializeVMCS(pCPU, (ULONG)GuestEsp)) ){
        KdPrint(("InitializeVMX Error!\n"));
        return STATUS_UNSUCCESSFUL;
    }
    
    KdPrint(("InitializeVMCS ok!\n"));
    
    //
    //  Launch It
    //
    
    Status = Virtualize( pCPU );
    if (Status == STATUS_UNSUCCESSFUL){
        KdPrint(("Virtualize Failed\n"));
        return STATUS_UNSUCCESSFUL;
    }
    return STATUS_SUCCESS;
}

NTSTATUS
DriverEntry(
    IN OUT PDRIVER_OBJECT   DriverObject,
    IN PUNICODE_STRING      RegistryPath
    )
{
    PDEVICE_OBJECT pdoDeviceObj = 0;
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    KIRQL OldIrql;
    ULONG i;
    
    pdoGlobalDrvObj = DriverObject;

    //
	// Create the device object.
    //
	
	if(!NT_SUCCESS(status = IoCreateDevice(
        DriverObject,
        0,
        &usDeviceName,
        FILE_DEVICE_UNKNOWN,
        FILE_DEVICE_SECURE_OPEN,
        FALSE,
        &pdoDeviceObj
        )))
    {
        //
		// Bail out (implicitly forces the driver to unload).
        //
		
		return status;
    }

    //
	// Now create the respective symbolic link object
    //
	
	if(!NT_SUCCESS(status = IoCreateSymbolicLink(
        &usSymlinkName,
        &usDeviceName
        )))
    {
        IoDeleteDevice(pdoDeviceObj);
        return status;
    }

    //
	// NOTE: You need not provide your own implementation for any major function that
    //       you do not want to handle. I have seen code using DDKWizard that left the
    //       *empty* dispatch routines intact. This is not necessary at all!
    //
	
	DriverObject->MajorFunction[IRP_MJ_CREATE] =
    DriverObject->MajorFunction[IRP_MJ_CLOSE] = VMXDBG_DispatchCreateClose;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = VMXDBG_DispatchDeviceControl;
    DriverObject->DriverUnload = VMXDBG_DriverUnload;
    
    //
    //  Is cpu support VT
    //
    
    if ( !CheckIfSupportVMX() ){
        KdPrint(("Do not support VT!\n"));
        return STATUS_SUCCESS;
    }
    
    //
    // Init GUI
    //

	GUIInitVideoInfo();

    for (i = 0; i < KeNumberProcessors; i++) {
        KeSetSystemAffinityThread((KAFFINITY) (1 << i));
        OldIrql = KeRaiseIrqlToDpcLevel();
        _StartVT();
        KeLowerIrql(OldIrql);
    }

    KdPrint(("All done...\n"));

    return STATUS_SUCCESS;
}
